package mj

import (
	"fmt"
	"github.com/go-redis/redis/v8"
	"gorm.io/gorm"
	"myai/core/types"
	"myai/service"
	"myai/service/oss"
	"myai/store"
	"myai/store/model"
	"myai/utils"
	"strings"
	"time"
)

// Service MJ 绘画服务
type Service struct {
	client          *Client // MJ Client
	taskQueue       *store.RedisQueue
	notifyQueue     *store.RedisQueue
	db              *gorm.DB
	wsService       *service.WebSocketService
	uploaderManager *oss.UploaderManager
	userService     *service.UserService
	clientIds       map[uint]string
}

func NewService(redisCli *redis.Client, db *gorm.DB, client *Client,
	manager *oss.UploaderManager, wsService *service.WebSocketService,
	userService *service.UserService) *Service {
	return &Service{
		db:              db,
		taskQueue:       store.NewRedisQueue("My_MidJourney_Task_Queue", redisCli),
		notifyQueue:     store.NewRedisQueue("My_MidJourney_Notify_Queue", redisCli),
		client:          client,
		wsService:       wsService,
		uploaderManager: manager,
		clientIds:       map[uint]string{},
		userService:     userService,
	}
}

type CBReq struct {
	Id          string      `json:"id"`
	Action      string      `json:"action"`
	Status      string      `json:"status"`
	Prompt      string      `json:"prompt"`
	PromptEn    string      `json:"promptEn"`
	Description string      `json:"description"`
	SubmitTime  int64       `json:"submitTime"`
	StartTime   int64       `json:"startTime"`
	FinishTime  int64       `json:"finishTime"`
	Progress    string      `json:"progress"`
	ImageUrl    string      `json:"imageUrl"`
	FailReason  interface{} `json:"failReason"`
	Properties  struct {
		FinalPrompt string `json:"finalPrompt"`
	} `json:"properties"`
}

// GetImageHash 获取图片hash
func GetImageHash(action string) string {
	// 从动作标识中获取 "MJ::JOB::upsample::1::26d56b92-4ed6-45e2-8561-563797923135"
	split := strings.Split(action, "::")
	if len(split) > 5 {
		return split[4]
	}
	return split[len(split)-1]
}

func (s *Service) Run() {
	// 将数据库中未提交的任务加载到队列中
	var jobs []model.MidJourneyJob
	s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
	for _, v := range jobs {
		var task types.MjTask
		err := utils.JsonDecode(v.TaskInfo, &task)
		if err != nil {
			log.Errorf("解码任务信息时出错：%v", err)
			continue
		}
		task.Id = v.Id
		s.clientIds[task.Id] = task.ClientId
		s.PushTask(task)
	}
	log.Infof("启动MidJourney作业消费者以获取服务")
	go func() {
		for {
			var task types.MjTask
			// 从队列中获取任务
			err := s.taskQueue.LPop(&task)
			if err != nil {
				log.Errorf("从队列中消费任务失败：%v", err)
				continue
			}

			// 判断提示词是否包含中文，如果有就进行翻译
			if utils.HasChinese(task.Prompt) {
				// 请求 大模型 对提示词进行翻译
				content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId)
				if err == nil {
					task.Prompt = content
				} else {
					log.Warnf("翻译提示词出错：%v", err)
				}
			}
			// 翻译负面提示词
			if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
				// 请求 大模型 对提示词进行翻译
				content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt), task.TranslateModelId)
				if err == nil {
					task.NegPrompt = content
				} else {
					log.Warnf("翻译提示词出错：%v", err)
				}
			}
			// 默认使用快速模式
			if task.Mode == "" {
				task.Mode = "fast"
			}

			s.clientIds[task.Id] = task.ClientId

			var job model.MidJourneyJob
			// 根据Id查询任务
			tx := s.db.Where("id = ?", task.Id).First(&job)
			if tx.Error != nil {
				log.Errorf("任务不存在，任务Id：%d", task.Id)
				continue
			}
			log.Infof("处理新的MidJourney任务：%+v", task)
			var res ImageRes
			switch task.Type {
			case types.TaskImage:
				res, err = s.client.Imagine(task)
				break
			case types.TaskUpscale:
				res, err = s.client.Upscale(task)
				break
			case types.TaskVariation:
				res, err = s.client.Variation(task)
				break
			case types.TaskBlend:
				res, err = s.client.Blend(task)
				break
			case types.TaskSwapFace:
				res, err = s.client.SwapFace(task)
				break
			}

			// 任务出现错误的处理
			if err != nil || (res.Code != 1 && res.Code != 22) {
				var errMsg string
				if err != nil {
					errMsg = err.Error()
				} else {
					errMsg = fmt.Sprintf("%v,%s", err, res.Description)
				}

				log.Errorf("绘画任务执行失败：%s", errMsg)
				job.Progress = service.FailTaskProgress
				job.ErrMsg = errMsg
				// 更新任务信息
				s.db.Updates(&job)
				// 任务失败，通知前端
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: task.ClientId,
					UserId:   task.UserId,
					JobId:    int(job.Id),
					Message:  service.TaskStatusFailed,
				})
				continue
			}
			// 任务成功的处理
			log.Infof("任务提交成功：%+v", res)
			// 更新任务 Id/频道
			job.TaskId = res.Result
			job.MessageId = res.Result
			job.ChannelId = res.Channel
			s.db.Updates(&job)
		}
	}()
}

// CheckTaskNotify 检查消息通知队列任务
func (s *Service) CheckTaskNotify() {
	go func() {
		for {
			var message service.NotifyMessage
			// 从队列中消费消息通知任务
			err := s.notifyQueue.LPop(&message)
			if err != nil {
				continue
			}
			log.Debugf("接收新的mj通知消息：%+v", message)
			// 获取一个客户端连接
			client := s.wsService.Clients.Get(message.ClientId)
			if client == nil {
				continue
			}
			// 发送消息
			utils.SendChannelMsg(client, types.ChMj, message.Message)
		}
	}()
}

// DownloadImages 把任务已经完成的图片进行下载
func (s *Service) DownloadImages() {
	go func() {
		var items []model.MidJourneyJob
		for {
			// 查询图片地址为空，并且进度为100的任务
			res := s.db.Where("img_url = ? and progress = ?", "", 100).Find(&items)
			if res.Error != nil {
				continue
			}
			// 下载图片
			for _, v := range items {
				if v.OrgURL == "" {
					continue
				}
				log.Infof("尝试下载图片：%s", v.OrgURL)
				// 如果返回的时 discord 图片地址，则使用代理下载
				proxy := false
				if strings.HasPrefix(v.OrgURL, "https://cdn.discordapp.com") {
					proxy = true
				}

				// 下载图片
				imgURL, err := s.uploaderManager.GetUploadHandler().PutUrlFile(v.OrgURL, proxy)
				if err != nil {
					log.Errorf("下载图片失败，地址：%s，错误：%v", v.OrgURL, err)
					continue
				} else {
					log.Infof("成功下载图片：%s", imgURL)
				}

				v.ImgURL = imgURL
				// 更新任务图片地址
				s.db.Updates(&v)
				// 推送通知消息
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: s.clientIds[v.Id],
					UserId:   int(v.UserId),
					JobId:    int(v.Id),
					Message:  service.TaskStatusFinished,
				})
			}
			// 每隔五秒执行一次
			time.Sleep(time.Second * 5)
		}
	}()
}

// PushTask 将新的mj任务推入任务队列
func (s *Service) PushTask(task types.MjTask) {
	log.Debugf("将新的MidJourney任务添加到任务列表中：%+v", task)
	s.taskQueue.RPush(task)
}

// SyncTaskProgress 异步拉取任务
func (s *Service) SyncTaskProgress() {
	go func() {
		var jobs []model.MidJourneyJob
		for {
			// 查询 MJ 任务进度小于100的
			err := s.db.Where("progress < ?", 100).Find(&jobs).Error
			if err != nil {
				continue
			}
			for _, job := range jobs {
				// 判断任务 10 分钟还没完成标记为失败
				if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
					job.Progress = service.FailTaskProgress
					job.ErrMsg = "任务超时"
					// 更新任务信息
					s.db.Updates(&job)
					continue
				}

				if job.ChannelId == "" {
					continue
				}
				// 查询任务
				task, err := s.client.QueryTask(job.TaskId, job.ChannelId)
				if err != nil {
					log.Errorf("查询任务失败：%v", err)
					continue
				}
				// 任务执行失败了
				if task.FailReason != "" {
					s.db.Model(&model.MidJourneyJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
						"progress": service.FailTaskProgress,
						"err_msg":  task.FailReason,
					})
					log.Errorf("任务执行失败：%v", task.FailReason)
					// 消息通知
					s.notifyQueue.RPush(service.NotifyMessage{
						ClientId: s.clientIds[job.Id],
						UserId:   int(job.UserId),
						JobId:    int(job.Id),
						Message:  service.TaskStatusFailed,
					})
					continue
				}

				// 判断是否有可执行按钮（动作）
				if len(task.Buttons) > 0 {
					// 从 CustomId（动作标识）中截取hash
					job.Hash = GetImageHash(task.Buttons[0].CustomId)
				}

				oldProgress := job.Progress
				// 获取新的任务进度
				job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
				// 获取图片地址
				if task.ImageUrl != "" {
					job.OrgURL = task.ImageUrl
				}
				// 更新任务信息
				err = s.db.Updates(&job).Error
				if err != nil {
					log.Errorf("任务信息更新到数据库失败：%v", err)
					continue
				}
				// 通知前端更新任务进度
				if oldProgress != job.Progress {
					message := service.TaskStatusRunning
					if job.Progress == 100 {
						message = service.TaskStatusFinished
					}
					// 发送消息通知
					s.notifyQueue.RPush(service.NotifyMessage{
						ClientId: s.clientIds[job.Id],
						UserId:   int(job.UserId),
						JobId:    int(job.Id),
						Message:  message,
					})
				}
			}
			// 找出失败的任务，并恢复其扣减的算力
			s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
			for _, job := range jobs {
				// 增加算力
				err = s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  "mid-journey",
					Remark: fmt.Sprintf("任务失败，退回算力。任务Id：%d，错误信息：%s", job.Id, job.ErrMsg),
				})
				if err != nil {
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 5)
		}
	}()
}
